TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
TokenGate：可学习 token 剪枝加速 LLM 预填充且精度不降
🧭 LLM 长上下文预填充加速 · 可学习 token 剪枝 ★ 深度解读
Stanford University · UC Berkeley
提出可插入组件 TokenGate，通过类 MoE 的 token-gated attention 端到端训练 token 剪枝，大幅加速长上下文预填充且精度不降
📌 工作参考价值（4 点）
- 在电商/Agent 场景中，长文档（商品详情、用户评论、搜索广告上下文、RAG 检索文档）的预填充是延迟和内存瓶颈；可借鉴 TokenGate 在模型中间层引入可学习 token 剪枝，相比 FastKV 等启发式方法，端到端训练能显著提升速度且精度无损，且可与 LoRA 一起微调，训练开销仅约 1.3x/step，总体甚至更少。
- token-gated attention 是关键 trick：把不可微的 top-K 剪枝通过类似 MoE 的门控值乘到 attention logits 上，让 scoring network 获得梯度，同时让模型骨干自适应剪枝后的稀疏注意力；该思路可迁移到任何需要在 Transformer 中间层做软筛选的任务，如多文档 RAG 中对检索片段做软裁剪。
- 训练时使用 dual budget sampling（全量 + 对数均匀采样的剪枝预算各一半损失），使单个模型泛化到不同 token 预算，无需为每个压缩比重新训练；这对线上需要动态调整 TTFT 和内存的推荐 / Agent 系统非常实用，可以一次训练，推理时按需调整 K。
- 工程实现上，将 per-key gate bias 融合进 FlashAttention 内核、剪枝后 forward 用 CUDA graph 固定 padding、KV cache 选择性裁剪等做法可直接复用到自研推理引擎，降低集成成本。